package com.swak.job.worker;

import java.util.List;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionStage;
import java.util.function.Function;

import com.swak.Constants;
import com.swak.job.Job;
import com.swak.job.JobConfig;
import com.swak.job.JobDispatch;
import com.swak.job.JobFragments;
import com.swak.job.JobTask;
import com.swak.job.enums.Role;
import com.swak.job.exception.JobException;
import com.swak.job.queue.QueueService;
import com.swak.job.storage.StorageService;
import com.swak.job.timer.Timer;
import com.swak.job.timer.Timer.CronTrigger;
import com.swak.job.timer.TimerTask.TaskFrag;
import com.swak.leader.LeaderService;
import com.swak.loadbalance.LoadBalance;
import com.swak.loadbalance.WeightAble;
import com.swak.meters.MetricsFactory;
import com.swak.registry.NotifyListener;
import com.swak.registry.RegistryService;
import com.swak.registry.URL;
import com.swak.utils.Ints;
import com.swak.utils.Lists;
import com.swak.utils.time.DateTimes;

import lombok.Getter;
import lombok.Setter;
import lombok.experimental.Accessors;

/**
 * 任务分发：主节点才能分发任务
 * 
 * @author lifeng
 * @date 2021年2月25日 下午2:30:20
 */
public class DispatchWorker extends LeaderAbstractWorker implements NotifyListener {

	protected final StorageService storageService;

	private Executers executers;
	private int dispatchTasks;

	private String cronExpression;
	private LoadBalance<Executer> loadBalance;
	private URL subscribeURL;

	/**
	 * 执行器
	 */
	private Timer timer;
	private CronTrigger cronTrigger;

	/**
	 * 处理消费请求
	 */
	Function<JobTask, Object> jobTaskConsumeFunc = this::handleJobTaskRespone;

	public DispatchWorker(Timer timer, JobConfig config, Job job, StorageService storageService,
			QueueService queueService, RegistryService registryService, LeaderService leaderService,
			MetricsFactory metricsFactory) {
		super(config, Role.Dispatcher, job, queueService, registryService, leaderService, metricsFactory);
		this.storageService = storageService;
		this.timer = timer;
		this.dispatchTasks = this.job.getDispatchTasks();
	}

	@Override
	protected void doRefresh() {
		super.doRefresh();
		this.job.setOrder(this.getRefreshAbleOrder());
		this.cronExpression = this.getRefreshAbleCronExpression();
		this.loadBalance = this.getRefreshAbleLoadBalances().newLoadBalance();
		this.subscribeURL = URL.valueOf(this.url).addParameter(Constants.Category_Key, Role.Executer.name());
		if (this.executers != null) {
			this.loadBalance.onRefresh(this.executers.getOpenedExecuters());
		}
	}

	@Override
	public void start() {
		super.start();
		this.queueService.createQueue(queue);
		this.queueService.consumeQueue(queue, jobTaskConsumeFunc);
		this.storageService.create(job);
		this.registryService.subscribe(subscribeURL, this);
		this.cronTrigger = this.timer.scheduleTasks(job.getName(), this.cronExpression, job.isRunOnBoot(),
				this::doTask);
	}

	private Object handleJobTaskRespone(JobTask task) {
		if (logger.isDebugEnabled()) {
			logger.debug("Job {} Execute Task {} Finish.", this.job.getName(), task.getConfig());
		}
		this.storageService.execute(task);
		return task;
	}

	private Object doTask(final TaskFrag frag) {
		if (this.opened && this.isLeader()) {
			return this.doDispatch(frag);
		}
		return null;
	}

	@Override
	public void notify(List<URL> workers) {
		List<Executer> es = Lists.newArrayList();
		for (URL url : workers) {
			es.add(new Executer(url));
		}
		this.executers = new Executers(es);
		this.loadBalance.onRefresh(this.executers.getOpenedExecuters());
	}

	protected CompletableFuture<Void> doDispatch(TaskFrag frag) {
		try {

			// 一次调度
			JobDispatch dispatch = job.newDispatch().setUrl(this.url.toFullString()).setCreateTime(frag.getRunTime());

			// 表示此次调度成功
			if (this.storageService.dispatch(dispatch, frag)) {
				JobFragments fragments = new JobFragments(job, dispatch, executers.getWeights());
				Object result = null;
				if (this.method != null) {
					Object[] args = new Object[] { fragments };
					result = this.invoke(args);
				} else {
					for (int i = 1; i <= this.dispatchTasks && i <= fragments.getFragments(); i++) {
						fragments.addConfig("" + i);
					}
				}
				return this.dispatchTask(result, fragments);
			}

			// 不需要处理
			return CompletableFuture.completedFuture(null);
		} catch (Throwable e) {
			throw new JobException(e);
		}
	}

	/**
	 * 分发任务
	 * 
	 * @param result
	 */
	@SuppressWarnings("unchecked")
	protected CompletableFuture<Void> dispatchTask(Object result, JobFragments fragments) {
		CompletableFuture<Void> future = new CompletableFuture<>();
		if (result != null && result instanceof CompletionStage) {
			CompletionStage<Object> resultFuture = (CompletionStage<Object>) result;
			resultFuture.whenComplete((r, e) -> {
				this.dispatchTasks(future, fragments, e);
			});
		} else {
			this.dispatchTasks(future, fragments, null);
		}
		return future;
	}

	protected void dispatchTasks(CompletableFuture<Void> future, JobFragments fragments, Throwable e) {
		try {
			if (e != null) {
				future.completeExceptionally(e);
				return;
			}

			// 任务分片配置
			List<String> configs = fragments.getConfigs();
			if (logger.isDebugEnabled()) {
				logger.debug("Job {} Dispatch {} Tasks.", this.job.getName(), configs.size());
			}
			for (String config : configs) {
				this.doWithMetrics(fragments, config);
			}
			future.complete(null);
		} catch (Exception ex) {
			future.completeExceptionally(ex);
		}
	}

	protected void doWithMetrics(JobFragments fragments, String config) {
		Object metrics = this.namedMetrics != null ? this.namedMetrics.begin() : null;
		boolean success = true;
		try {
			Executer executer = this.loadBalance.select();
			if (executer == null) {
				logger.error("Job {} Need Opened Executer.", this.job.getName());
				return;
			}
			JobTask task = new JobTask().setExecuteUrl(executer.url.toFullString())
					.setDispatchUrl(this.url.toFullString()).setDispatchId(fragments.getDispatch().getId())
					.setJobId(job.getId()).setJobName(job.getName()).setConfig(config).setCreateTime(DateTimes.now())
					.setRetryTimes(job.getRetryTimes());
			this.storageService.createTask(task);
			this.dispatchTask(executer, task);
			this.storageService.dispatchTask(task);
		} catch (Exception e) {
			success = false;
		} finally {
			if (this.namedMetrics != null) {
				this.namedMetrics.end(metrics, success);
			}
		}
	}

	// 生成一次任务
	protected void dispatchTask(Executer executer, JobTask task) {
		try {
			this.queueService.publishQueue(executer.queue, task);
		} catch (Exception e) {
			logger.error("Job {} Dispatch Task {} Error.", job.getName(), task.getId(), e);
		}
	}

	@Override
	public void stop() {
		super.stop();
		if (this.cronTrigger != null) {
			this.cronTrigger.cancel();
		}
		this.registryService.unsubscribe(subscribeURL, this);
		this.queueService.unconsumeQueue(queue, jobTaskConsumeFunc);
	}

	/**
	 * 可用的执行器
	 * 
	 * @author DELL
	 */
	@Getter
	@Setter
	@Accessors(chain = true)
	class Executers {
		private List<Executer> executers;
		private Integer weights;

		Executers(List<Executer> executers) {
			this.executers = executers;
			for (Executer executer : executers) {
				if (executer.isOpen()) {
					weights = Ints.addI(weights, executer.getWeight());
				}
			}
		}

		List<Executer> getOpenedExecuters() {
			List<Executer> openeds = Lists.newArrayList();
			for (Executer executer : executers) {
				if (executer.isOpen()) {
					openeds.add(executer);
				}
			}
			return openeds;
		}

	}

	/**
	 * 执行器
	 * 
	 * @author lifeng
	 */
	@Getter
	@Setter
	@Accessors(chain = true)
	class Executer implements WeightAble {
		private final URL url;
		private final String queue;
		private final int weight;
		private final boolean open;

		Executer(URL url) {
			this.url = url;
			this.queue = queueService.parseExecuteQueue(url);
			this.weight = url.getParameter(Weight_key, 1);
			this.open = url.getParameter(Open_key, true);
		}
	}
}
